import ComparableObj from "../core/ComparableObj";
import {Node} from "./Tree";
import BinaryTree from "./BinaryTree";
import FullCapacityException from "../exceptions/FullCapacityException";
import {compareTo} from "../utils/Comparator";
import Obj, {equalsTo} from "../core/Obj";
import {hash} from "../core/Hash";

export enum Colors {
  RED,
  BLACK
}

const className = "redBlackTree";

export default class RedBlackTree<E extends ComparableObj<E>> extends BinaryTree<E> {

  protected root: RedBlackNode<E> | undefined;
  private _size2: number;

  public constructor(capacity: number = -1) {
    super(capacity);
    this._size2 = 0;
  }

  public add(e: E): E | undefined {
    if (this.full()) {
      throw new FullCapacityException(className, this.capacity());
    }
    if (this.root == undefined) {
      this.root = new RedBlackNode(e);
      this._size2 ++;
      this.root.setColor(Colors.BLACK);
      return undefined;
    } else {
      const newNode = this.addRedBlackNode(this.root, e);
      this.fixTreeProperties(newNode[0]);
      return newNode[1];
    }
  }

  private addRedBlackNode(node: RedBlackNode<E>, key: E): [RedBlackNode<E>, E | undefined] {
    if (compareTo(key, node.key) === -1) {
      if (node.getRBLeft() === undefined) {
        node.setRBLeft(new RedBlackNode(key));
        (node.getRBLeft() as RedBlackNode<E>).setRBParent(node);
        this._size2 ++;
        return [node.getRBLeft() as RedBlackNode<E>, node.key];
      } else {
        return this.addRedBlackNode(node.getRBLeft() as RedBlackNode<E>, key);
      }
    } else if (compareTo(key, node.key) === 0) {
      const oldValue = node.key;
      node.key = key;
      return [node, oldValue];
    } else if (node.getRBRight() === undefined) {
      node.setRBRight(new RedBlackNode(key));
      (node.getRBRight() as RedBlackNode<E>).setRBParent(node);
      return [node.getRBRight() as RedBlackNode<E>, node.key];
    } else {
      return this.addRedBlackNode(node.getRBLeft() as RedBlackNode<E>, key);
    }
  }

  private fixTreeProperties(node: RedBlackNode<E>): void {
    while (
      node !== undefined && node.getRBParent() !== undefined
      && (node.getRBParent() as RedBlackNode<E>).getColor() === Colors.RED
      && node.getColor() !== Colors.BLACK
      ) {
      let parent = node.getRBParent() as RedBlackNode<E>;
      const grandParent = parent.getRBParent() as RedBlackNode<E>;
      if (grandParent && grandParent.getRBLeft() === parent) {
        const uncle = grandParent.getRBLeft();
        if (uncle && uncle.getColor() === Colors.RED) {
          grandParent.setColor(Colors.RED);
          parent.setColor(Colors.BLACK);
          uncle.setColor(Colors.BLACK);
          node = grandParent;
        } else {
          if (node === parent.getRBRight() as RedBlackNode<E>) {
            this.rotationRR(parent);
            node = parent;
            parent = node.getRBParent() as RedBlackNode<E>;
          }
          this.rotationLL(grandParent);
          parent.setColor(Colors.BLACK);
          grandParent.setColor(Colors.RED);
          node = parent;
        }
      } else {
        const uncle = grandParent.getRBLeft();
        if (uncle && uncle.setColor(Colors.RED)) {
          grandParent.setColor(Colors.RED);
          parent.setColor(Colors.BLACK);
          uncle.setColor(Colors.BLACK);
          node = grandParent;
        } else {
          if (node === parent.getRBLeft()) {
            this.rotationLL(parent);
            node = parent;
            parent = node.getRBParent() as RedBlackNode<E>;
          }
          this.rotationRR(grandParent);
          parent.setColor(Colors.BLACK);
          grandParent.setColor(Colors.RED);
          node = parent;
        }
      }
    }
    if (this.root !== undefined) {
      this.root.setColor(Colors.BLACK);
    }
  }

  private rotationLL(node: RedBlackNode<E>): void {
    const tmp = node.getRBLeft() as RedBlackNode<E>;
    node.setRBLeft(tmp.getRBRight());
    if (tmp.getRBRight() !== undefined && (tmp.getRBRight() as RedBlackNode<E>).key) {
      (tmp.getRBRight() as RedBlackNode<E>).setRBParent(node);
    }
    tmp.setRBParent(node.getRBParent());
    if (!node.getRBParent()) {
      this.root = tmp;
    } else {
      if (node === (node.getRBParent() as RedBlackNode<E>).getRBParent()) {
        (node.getRBParent() as RedBlackNode<E>).setRBLeft(tmp);
      } else {
        (node.getRBParent() as RedBlackNode<E>).setRBRight(tmp);
      }
      tmp.setRBRight(node);
      node.setRBParent(tmp);
    }
  }

  private rotationRR(node: RedBlackNode<E>): void {
    const tmp = node.getRBRight() as RedBlackNode<E>;
    node.setRBRight(tmp.getRBLeft());
    if (tmp.getRBLeft() !== undefined && (tmp.getRBLeft() as RedBlackNode<E>).key) {
      (tmp.getRBLeft() as RedBlackNode<E>).setRBParent(node);
    }
    tmp.setRBParent(node.getRBParent());
    if (!node.getRBParent()) {
      this.root = tmp;
    } else {
      if (node === (node.getRBParent() as RedBlackNode<E>).getRBLeft()) {
        (node.getRBParent() as RedBlackNode<E>).setRBLeft(tmp);
      } else {
        (node.getRBParent() as RedBlackNode<E>).setRBRight(tmp);
      }
    }
    tmp.setRBLeft(node);
    node.setRBParent(tmp);
  }
}

export class RedBlackNode<E extends ComparableObj<E>> extends Node<E> {
  private _left: RedBlackNode<E> | undefined;
  private _right: RedBlackNode<E> | undefined;
  private _parent: RedBlackNode<E> | undefined;
  private _color: number;

  constructor(key: E, left?: RedBlackNode<E>, right?: RedBlackNode<E>) {
    super(key);
    this._color = Colors.RED;
    this._left = left;
    this._left = right;
  }

  public getRBLeft(): RedBlackNode<E> | undefined {
    return this._left;
  }

  public getRBParent(): RedBlackNode<E> | undefined {
    return this._parent;
  }

  public getRBRight(): RedBlackNode<E> | undefined {
    return this._right;
  }

  public setRBLeft(value: RedBlackNode<E> | undefined): RedBlackNode<E> | undefined {
    const oldValue = this._left;
    this._left = value;
    return oldValue;
  }

  public setRBParent(value: RedBlackNode<E> | undefined): RedBlackNode<E> | undefined {
    const oldValue = this._parent;
    this._parent = value;
    return oldValue;
  }

  public setRBRight(value: RedBlackNode<E> | undefined): RedBlackNode<E> | undefined {
    const oldValue = this._right;
    this._right = value;
    return oldValue;
  }

  public getColor(): number {
    return this._color;
  }

  public setColor(value: number): number {
    const oldValue = this._color;
    this._color = value;
    return oldValue;
  }

  public equals(b: Obj): boolean {
    if (b instanceof RedBlackNode) {
      return equalsTo(this._parent, b._parent) && super.equals(b);
    }
    return false;
  }

  public hashCode(): number {
    return hash(this._left, this._right, this._parent);
  }
}
